import torch
import torch.nn.functional as F
from torch.optim import Adam
from torch_geometric.data import DataLoader, DenseDataLoader as DenseLoader
from torch import tensor
import torch_geometric.transforms as T
from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR
import pdb
import random
import numpy as np
from torch.autograd import grad
from torch_geometric.data import Batch
from utils import k_fold, num_graphs
from tqdm import tqdm

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# process_dataset 函数用于处理数据集。它首先统计数据集中的节点数，然后将数据集中的每个图转换为稠密表示（dense representation）。
# 具体来说，它将每个图的节点特征矩阵扩展为固定大小，不足部分填充为零，并将所有图处理为相同的节点数。这是因为某些图神经网络模型（如GCN）需要固定大小的输入特征
def process_dataset(dataset):
    
    num_nodes = max_num_nodes = 0
    for data in dataset:
        num_nodes += data.num_nodes
        max_num_nodes = max(data.num_nodes, max_num_nodes)
    # 这行代码将数据集中所有图的平均节点数扩展为 num_nodes，扩展的数量是当前数据集平均节点数的5倍，但不超过数据集中的最大节点数。这样做是为了将所有图的节点数调整为一个固定的值，以便后续的处理
    num_nodes = min(int(num_nodes / len(dataset) * 5), max_num_nodes)
    # 这行代码创建了一个将图转换为稠密表示的转换器。T.ToDense(num_nodes) 表示将图的节点特征转换为稠密张量，将每个图的节点特征向量扩展到 num_nodes 的长度，并且用零填充不足的部分
    transform = T.ToDense(num_nodes)
    new_dataset = []
    # 使用 tqdm 函数创建一个进度条，用于迭代数据集中的每个图。tqdm 可以显示迭代的进度，让用户知道程序的执行进展
    for data in tqdm(dataset):
        # 对每个图应用上面定义的转换器 transform，将图的节点特征转换为稠密表示
        data = transform(data)
        # 计算需要添加的零填充的数量，即 num_nodes 和当前图节点数的差值
        add_zeros = num_nodes - data.feat.shape[0]
        # 检查是否需要进行填充操作，即判断 add_zeros 是否大于0
        if add_zeros:
            # 如果需要填充，则创建一个全零张量，其大小为 (add_zeros, dim)，其中 dim 表示节点特征的维度。然后，使用 torch.cat 函数将零张量与当前图的节点特征张量在垂直方向进行拼接
            dim = data.feat.shape[1]
            data.feat = torch.cat((data.feat, torch.zeros(add_zeros, dim)), dim=0)
        # 将处理后的图数据添加到 new_dataset 列表中
        new_dataset.append(data)
    return new_dataset

# train_baseline_syn 函数实现了基线模型的训练过程。
# 它首先创建了用于训练、验证和测试的数据加载器。
# 然后初始化了模型、优化器和学习率调度器。
# 在每个训练周期中，通过调用 train 函数进行模型训练，并调用 eval_acc 函数评估模型在验证集和测试集上的准确率。
# 同时，使用余弦退火学习率调度器来动态调整学习率。最后，打印并记录了训练过程中的关键指标

# 这是一个函数定义，名为 train_baseline_syn，用于训练基线模型。它接收四个参数：train_set、val_set 和 test_set 分别表示训练集、验证集和测试集，model_func 是一个模型构建函数，args 是用于配置模型训练的参数
def train_baseline_syn(train_set, val_set, test_set, model_func=None, args=None):

    # 这三行代码分别使用 PyTorch Geometric 提供的 DataLoader 类创建训练、验证和测试集的数据加载器。数据加载器用于批量加载数据，并支持自动分批、打乱数据和并行加载等功能
    train_loader = DataLoader(train_set, args.batch_size, shuffle=True)
    val_loader = DataLoader(val_set, args.batch_size, shuffle=False)
    test_loader = DataLoader(test_set, args.batch_size, shuffle=False)

    # 这行代码检查 args 参数中的 feature_dim 是否为 -1，如果是，则将其设置为 args 参数中的 max_degree
    if args.feature_dim == -1:
        args.feature_dim = args.max_degree
    # 这行代码根据传入的 model_func 构建模型，并将模型移动到可用的设备（GPU 或 CPU）上
    model = model_func(args.feature_dim, args.num_classes).to(device)
    # 这行代码使用 Adam 优化器来优化模型的参数，学习率为 args.lr。
    optimizer = Adam(model.parameters(), lr=args.lr)

    # 这行代码创建一个余弦退火学习率调度器 (CosineAnnealingLR)，用于动态调整学习率。T_max 参数指定了一次学习率周期的迭代次数，eta_min 参数指定了学习率的最小值，last_epoch 参数指定了初始的迭代次数。
    lr_scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs, eta_min=args.min_lr, last_epoch=-1, verbose=False)
    # 初始化四个变量 best_val_acc、update_test_acc、update_train_acc 和 update_epoch，分别用于记录最佳验证准确率、相应的测试准确率、相应的训练准确率和对应的迭代次数
    best_val_acc, update_test_acc, update_train_acc, update_epoch = 0, 0, 0, 0

    # 使用循环迭代训练多个周期，每个周期的次数由 args.epochs 决定
    for epoch in range(1, args.epochs + 1):

        # 调用 train() 函数进行模型训练，返回训练损失和训练准确率；调用 eval_acc() 函数评估模型在验证集和测试集上的准确率
        train_loss, train_acc = train(model, optimizer, train_loader, device, args)
        val_acc = eval_acc(model, val_loader, device, args)
        test_acc = eval_acc(model, test_loader, device, args)
        # 使用学习率调度器进行学习率调整。
        lr_scheduler.step()
        # 如果当前验证准确率超过了历史最佳验证准确率，则更新历史最佳验证准确率和对应的测试准确率、训练准确率以及迭代次数
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            update_test_acc = test_acc
            update_train_acc = train_acc
            update_epoch = epoch

        # 打印当前周期的训练损失、训练准确率、验证准确率和测试准确率等信息，并显示最佳验证准确率和对应的测试准确率的更新情况
        print("BIAS:[{:.2f}] | Model:[{}] Epoch:[{}/{}] Loss:[{:.4f}] Train:[{:.2f}] val:[{:.2f}] Test:[{:.2f}] | Best Val:[{:.2f}] Update Test:[{:.2f}] at Epoch:[{}] | lr:{:.6f}"
                .format(args.bias,
                        args.model,
                        epoch, 
                        args.epochs,
                        train_loss, 
                        train_acc * 100, 
                        val_acc * 100,
                        test_acc * 100, 
                        best_val_acc * 100,
                        update_test_acc * 100, 
                        update_epoch,
                        optimizer.param_groups[0]['lr']))

    # 打印最终的最佳验证准确率和对应的测试准确率，以及它们对应的迭代周期。
    print("syd: BIAS:[{:.2f}] | Best Val acc:[{:.2f}] Test acc:[{:.2f}] at epoch:[{}]"
        .format(args.bias,
                best_val_acc * 100, 
                update_test_acc * 100,
                update_epoch))

# num_graphs 函数用于计算数据集中的图数量
def num_graphs(data):
    if data.batch is not None:
        return data.num_graphs
    else:
        x = data.x if data.x is not None else data.feat
        return x.size(0)

# train 函数用于模型的训练过程。它首先将模型设置为训练模式，然后遍历数据加载器中的每个批次，执行前向传播、计算损失、反向传播和参数更新操作。最后返回平均损失和准确率
def train(model, optimizer, loader, device, args):

    # 将模型设置为训练模式，这将启用模型中的训练特定操作，如 dropout 和 batch normalization。
    model.train()
    # 初始化变量 total_loss 和 correct，用于分别跟踪总损失和正确预测的数量
    total_loss = 0
    correct = 0

    # 使用 enumerate() 函数迭代加载器中的数据批次，并获得每个批次的索引 it 和数据 data
    for it, data in enumerate(loader):
        # 在每个批次开始时，清除之前批次的梯度信息，以确保梯度不会在多个批次之间累积
        optimizer.zero_grad()
        # 将数据移动到指定的设备 (GPU 或 CPU) 上进行计算
        data = data.to(device)
        # 使用模型对数据进行前向传播，获得模型的预测输出 out
        out = model(data)
        # 计算模型的预测输出 out 与真实标签 data.y 之间的负对数似然损失。nll_loss 函数用于计算多分类问题中的负对数似然损失
        loss = F.nll_loss(out, data.y.view(-1))
        # 根据模型的预测输出 out，获取每个样本的预测类别，并将预测概率最大的类别作为模型的预测结果 pred。
        pred = out.max(1)[1]
        # 将模型的预测结果 pred 与真实标签 data.y 进行比较，并统计正确预测的样本数量，然后累加到 correct 变量中。
        correct += pred.eq(data.y.view(-1)).sum().item()
        # 对损失值 loss 进行反向传播，计算模型参数的梯度
        loss.backward()
        # 将当前批次的损失值乘以批次中的图的数量，并累加到 total_loss 变量中
        total_loss += loss.item() * num_graphs(data)
        # 使用优化器更新模型参数，通过使用之前计算的梯度值来沿着梯度方向更新参数
        optimizer.step()
    # 返回整个训练集上的平均损失值 (total_loss / len(loader.dataset)) 和训练准确率 (correct / len(loader.dataset))。
    return total_loss / len(loader.dataset), correct / len(loader.dataset)


# 这是一个名为 eval_acc 的函数，用于评估模型在验证集或测试集上的准确率。它接受模型 (model)、数据加载器 (loader)、设备 (device) 和参数 (args) 作为输入
def eval_acc(model, loader, device, args):
    # 将模型设置为评估模式，这将关闭模型中的训练特定操作，如 dropout 和 batch normalization。
    model.eval()
    # 初始化变量 correct，用于跟踪正确预测的数量。
    correct = 0
    # 迭代加载器中的数据批次，依次处理每个批次的数据
    for data in loader:
        # 将数据移动到指定的设备 (GPU 或 CPU) 上进行计算
        data = data.to(device)
        # 使用 torch.no_grad() 上下文管理器，确保在评估阶段不会计算梯度
        with torch.no_grad():
            # 使用模型对数据进行前向传播，获得模型的预测输出 pred。然后，取出每个样本预测的类别，并将概率最大的类别作为模型的预测结果。
            pred = model(data).max(1)[1]
        # 将模型的预测结果 pred 与真实标签 data.y 进行比较，并统计正确预测的样本数量，然后累加到 correct 变量中。
        correct += pred.eq(data.y.view(-1)).sum().item()
    # 返回整个验证集或测试集上的准确率，即正确预测的样本数量除以总样本数量 (correct / len(loader.dataset))
    return correct / len(loader.dataset)
